package bjut.zrpc.core.optimization.protect;

import bjut.zrpc.core.optimization.bean.ZRpcServiceBean;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;

/**
 * @Author: zhuzhiming
 * @Date: 2022/6/29 16:53
 */
public class LimitStream {

    private static final ConcurrentHashMap<String, Long> CONSUMER_LIMIT_MAP = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, Map<Long, Integer>> LIMIT_COUNT_MAP = new ConcurrentHashMap<>();
    public static final ConcurrentHashMap<String, ZRpcServiceBean> RPCSERVICE_MAP = new ConcurrentHashMap<>();

    public boolean needLimit(String serviceNameAndClientAddress, int limitCount, int inTime, int duration) {

        Long window = CONSUMER_LIMIT_MAP.get(serviceNameAndClientAddress);
        long currentTime = System.currentTimeMillis();

        if (window == null) {
            window = System.currentTimeMillis();
        }


        Map<Long, Integer> limitMap = LIMIT_COUNT_MAP.get(serviceNameAndClientAddress);

        if (limitMap == null) {
            HashMap<Long, Integer> map = new HashMap<>();
            map.put(window, 0);
            LIMIT_COUNT_MAP.put(serviceNameAndClientAddress, map);
            limitMap = LIMIT_COUNT_MAP.get(serviceNameAndClientAddress);
        }

        int currentCount = limitMap.get(window) + 1;

        CONSUMER_LIMIT_MAP.put(serviceNameAndClientAddress, window);
        LIMIT_COUNT_MAP.get(serviceNameAndClientAddress).put(window, currentCount);
        LIMIT_COUNT_MAP.put(serviceNameAndClientAddress, LIMIT_COUNT_MAP.get(serviceNameAndClientAddress));

        if ((currentTime > window + inTime && currentCount < limitCount) || currentTime >= window + duration) {
            CONSUMER_LIMIT_MAP.remove(serviceNameAndClientAddress);
            LIMIT_COUNT_MAP.remove(serviceNameAndClientAddress);
            return false;
        }

        return currentTime <= window + inTime && currentCount >= limitCount;

    }

    public ZRpcServiceBean getServiceBean(String serviceName){
        return RPCSERVICE_MAP.get(serviceName);
    }


}
